from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
import os
from mp_api.client import MPRester
from datetime import datetime
import json
import asyncio
from baidu_search import baidu_searcher
import httpx
import json as json_module

# 读取AI配置
def load_ai_config():
    """加载AI配置"""
    config_path = "ai_config.json"
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json_module.load(f)
        return config
    except FileNotFoundError:
        # 如果配置文件不存在，使用默认配置
        default_config = {
            "ai_api": {
                "url": "http://office20.inoteexpress.com:8003/v1/completions",
                "model": "S1-Base-32B",
                "default_temperature": 0.7,
                "default_max_tokens": 7000,
                "timeout": 60.0,
                "stop_tokens": ["<|im_end|>"]
            },
            "prompt_templates": {
                "system_prompt": "你是一个AI助手，精通材料科学和分子结构分析。请基于提供的材料数据，提供专业、准确的回答。",
                "material_info_template": "当前分析的材料信息：\n- 材料ID: {material_id}\n- 化学式: {formula_pretty}\n- 晶体系统: {crystal_system}\n- 空间群: {space_group}\n- 带隙: {band_gap} eV\n- 密度: {density} g/cm³\n- 体积: {volume} Å³"
            }
        }
        # 创建默认配置文件
        with open(config_path, 'w', encoding='utf-8') as f:
            json_module.dump(default_config, f, ensure_ascii=False, indent=2)
        return default_config

# 加载AI配置
AI_CONFIG = load_ai_config()

app = FastAPI(
    title="Molecule Scope API",
    description="API for accessing Materials Project data with 3D molecular visualization",
    version="1.0.0"
)

# 配置CORS中间件，允许Vue客户端访问
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:8002",
        "http://127.0.0.1:8002", 
        "http://localhost:8001",
        "http://127.0.0.1:8001",
        "http://office20.inoteexpress.com:8001",
        "http://office20.inoteexpress.com:8002",
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 添加OPTIONS请求处理
@app.options("/{path:path}")
async def options_handler(path: str):
    return {"message": "OK"}

# 添加自定义CORS头中间件
@app.middleware("http")
async def add_cors_headers(request, call_next):
    response = await call_next(request)
    # 确保所有响应都包含CORS头
    response.headers["Access-Control-Allow-Origin"] = "*"
    response.headers["Access-Control-Allow-Methods"] = "GET, POST, PUT, DELETE, OPTIONS"
    response.headers["Access-Control-Allow-Headers"] = "Content-Type, Authorization, X-Requested-With"
    response.headers["Access-Control-Allow-Credentials"] = "true"
    return response

# 从环境变量获取API密钥
API_KEY = os.getenv("MATERIALS_PROJECT_API_KEY",
                    "ETuMVaZhVYYQEEPWxsKbNn9j6nU1kPoD")


class MaterialSummary(BaseModel):
    material_id: str
    formula_pretty: str
    spacegroup: Optional[Dict[str, Any]] = None
    volume: Optional[float] = None
    density: Optional[float] = None
    elements: Optional[List[str]] = None
    band_gap: Optional[float] = None
    energy_per_atom: Optional[float] = None
    formation_energy_per_atom: Optional[float] = None
    total_magnetization: Optional[float] = None
    is_magnetic: Optional[bool] = None
    is_metal: Optional[bool] = None
    is_gap_direct: Optional[bool] = None
    efermi: Optional[float] = None
    xas: Optional[Dict[str, Any]] = None
    bulk_modulus: Optional[float] = None
    shear_modulus: Optional[float] = None
    energy_above_hull: Optional[float] = None
    surface_energy: Optional[float] = None
    work_function: Optional[float] = None
    universal_anisotropy: Optional[float] = None
    poisson_ratio: Optional[float] = None


class SearchRequest(BaseModel):
    query: str
    fields: Optional[List[str]] = None


# 常见材料的中英文名称映射
MATERIAL_NAME_MAPPING = {
    # 硅/Silicon
    "硅": "Si", "silicon": "Si", "silicon crystal": "Si",
    # 碳/钻石/Diamond
    "碳": "C", "carbon": "C", "钻石": "C", "diamond": "C", "graphite": "C",
    # 石英/二氧化硅/Quartz
    "石英": "SiO2", "二氧化硅": "SiO2", "quartz": "SiO2", "silica": "SiO2",
    # 铁/Iron
    "铁": "Fe", "iron": "Fe", "ferrum": "Fe",
    # 铝/Aluminum
    "铝": "Al", "aluminum": "Al", "aluminium": "Al",
    # 铜/Copper
    "铜": "Cu", "copper": "Cu",
    # 金/Gold
    "金": "Au", "gold": "Au",
    # 银/Silver
    "银": "Ag", "silver": "Ag",
    # 水/Water/Ice
    "水": "H2O", "water": "H2O", "冰": "H2O", "ice": "H2O",
    # 盐/Salt
    "盐": "NaCl", "salt": "NaCl", "氯化钠": "NaCl",
    # 氧化铝/Alumina
    "氧化铝": "Al2O3", "alumina": "Al2O3", "刚玉": "Al2O3",
    # 氧化铁/Iron Oxide
    "氧化铁": "Fe2O3", "iron oxide": "Fe2O3", "赤铁矿": "Fe2O3",
    # 氧化铜/Copper Oxide
    "氧化铜": "CuO", "copper oxide": "CuO",
    # 氧化锌/Zinc Oxide
    "氧化锌": "ZnO", "zinc oxide": "ZnO",
    # 氧化钛/Titanium Oxide
    "氧化钛": "TiO2", "titanium oxide": "TiO2", "二氧化钛": "TiO2",
    # 氧化镁/Magnesium Oxide
    "氧化镁": "MgO", "magnesium oxide": "MgO",
    # 氧化钙/Calcium Oxide
    "氧化钙": "CaO", "calcium oxide": "CaO", "生石灰": "CaO",
    # 氢氧化钠/Sodium Hydroxide
    "氢氧化钠": "NaOH", "sodium hydroxide": "NaOH", "烧碱": "NaOH",
    # 硫酸/Sulfuric Acid
    "硫酸": "H2SO4", "sulfuric acid": "H2SO4",
    # 盐酸/Hydrochloric Acid
    "盐酸": "HCl", "hydrochloric acid": "HCl",
    # 硝酸/Nitric Acid
    "硝酸": "HNO3", "nitric acid": "HNO3",
    # 氨/Ammonia
    "氨": "NH3", "ammonia": "NH3",
    # 甲烷/Methane
    "甲烷": "CH4", "methane": "CH4",
    # 乙醇/Ethanol
    "乙醇": "C2H5OH", "ethanol": "C2H5OH", "酒精": "C2H5OH",
    # 葡萄糖/Glucose
    "葡萄糖": "C6H12O6", "glucose": "C6H12O6",
    # 蔗糖/Sucrose
    "蔗糖": "C12H22O11", "sucrose": "C12H22O11", "砂糖": "C12H22O11",
}


class PaperResult(BaseModel):
    title: str
    url: str
    abstract: str
    timestamp: str
    source: str


class PriceResult(BaseModel):
    title: str
    url: str
    abstract: str
    timestamp: str
    source: str
    price: Optional[str] = None
    currency: Optional[str] = None
    supplier: Optional[str] = None
    quantity: Optional[str] = None


class CrystalStructure(BaseModel):
    material_id: str
    lattice: Dict[str, Any]
    sites: List[Dict[str, Any]]
    charge: Optional[float] = None


class AIRequest(BaseModel):
    prompt: str
    material_data: Optional[Dict[str, Any]] = None
    temperature: float = 0.7
    max_tokens: int = 7000


class AIStreamResponse(BaseModel):
    content: str
    is_final: bool = False


@app.get("/")
async def root():
    return {"message": "Molecule Scope API is running", "version": "1.0.0"}


@app.get("/health")
async def health_check():
    return {"status": "healthy", "timestamp": datetime.now().isoformat()}


@app.get("/materials/search")
async def search_materials(query: str, limit: int = 10):
    """
    搜索材料 - 支持化学式、中英文名称、Materials Project ID搜索
    改进的搜索逻辑：
    1. 如果是mp-开头的ID，直接使用Materials Project API搜索
    2. 如果是化学式，使用Materials Project API按化学式搜索
    3. 如果是中文/英文名称，先尝试本地映射，然后使用外部搜索获取化学式
    """
    try:
        search_query = query.strip()
        
        # 1. 检查是否是Materials Project ID格式 (mp-开头后跟数字)
        is_material_id = (search_query.startswith('mp-') and
                         len(search_query) > 3 and
                         any(c.isdigit() for c in search_query[3:]))
        
        if is_material_id:
            # 直接按材料ID搜索
            with MPRester(API_KEY) as mpr:
                results = mpr.materials.summary.search(
                    material_ids=[search_query],
                    fields=["material_id", "formula_pretty", "symmetry.symbol",
                            "volume", "density", "elements", "band_gap"]
                )
                
                materials = []
                for doc in results:
                    materials.append({
                        "material_id": doc.material_id,
                        "formula_pretty": doc.formula_pretty,
                        "symmetry": doc.symmetry.dict() if hasattr(doc, 'symmetry') and doc.symmetry else None,
                        "volume": doc.volume,
                        "density": doc.density,
                        "elements": doc.elements,
                        "band_gap": doc.band_gap
                    })

                return {
                    "results": materials[:limit],
                    "count": len(materials),
                    "original_query": query,
                    "search_type": "material_id"
                }
        
        # 2. 检查是否是化学式（包含字母和数字，可能有括号）
        is_chemical_formula = _is_chemical_formula(search_query)
        
        if is_chemical_formula:
            # 直接按化学式搜索
            with MPRester(API_KEY) as mpr:
                results = mpr.materials.summary.search(
                    formula=search_query,
                    fields=["material_id", "formula_pretty", "symmetry.symbol",
                            "volume", "density", "elements", "band_gap"]
                )
                
                materials = []
                for doc in results:
                    materials.append({
                        "material_id": doc.material_id,
                        "formula_pretty": doc.formula_pretty,
                        "symmetry": doc.symmetry.dict() if hasattr(doc, 'symmetry') and doc.symmetry else None,
                        "volume": doc.volume,
                        "density": doc.density,
                        "elements": doc.elements,
                        "band_gap": doc.band_gap
                    })

                return {
                    "results": materials[:limit],
                    "count": len(materials),
                    "original_query": query,
                    "search_type": "formula"
                }
        
        # 3. 处理材料名称查询（中文/英文名称）
        # 首先检查本地名称映射
        query_lower = search_query.lower()
        if query_lower in MATERIAL_NAME_MAPPING:
            formula = MATERIAL_NAME_MAPPING[query_lower]
            print(f"Using local mapping: '{query}' -> '{formula}'")
            
            with MPRester(API_KEY) as mpr:
                results = mpr.materials.summary.search(
                    formula=formula,
                    fields=["material_id", "formula_pretty", "symmetry.symbol",
                            "volume", "density", "elements", "band_gap"]
                )
                
                materials = []
                for doc in results:
                    materials.append({
                        "material_id": doc.material_id,
                        "formula_pretty": doc.formula_pretty,
                        "symmetry": doc.symmetry.dict() if hasattr(doc, 'symmetry') and doc.symmetry else None,
                        "volume": doc.volume,
                        "density": doc.density,
                        "elements": doc.elements,
                        "band_gap": doc.band_gap
                    })

                return {
                    "results": materials[:limit],
                    "count": len(materials),
                    "original_query": query,
                    "search_type": "name_mapping",
                    "formula_used": formula
                }
        
        # 4. 如果本地映射中没有，尝试使用外部搜索获取化学式
        try:
            # 使用百度搜索获取可能的化学式
            external_formula = await _search_formula_from_external(query)
            if external_formula:
                print(f"External search found formula: '{query}' -> '{external_formula}'")
                
                with MPRester(API_KEY) as mpr:
                    results = mpr.materials.summary.search(
                        formula=external_formula,
                        fields=["material_id", "formula_pretty", "symmetry.symbol",
                                "volume", "density", "elements", "band_gap"]
                    )
                    
                    materials = []
                    for doc in results:
                        materials.append({
                            "material_id": doc.material_id,
                            "formula_pretty": doc.formula_pretty,
                            "symmetry": doc.symmetry.dict() if hasattr(doc, 'symmetry') and doc.symmetry else None,
                            "volume": doc.volume,
                            "density": doc.density,
                            "elements": doc.elements,
                            "band_gap": doc.band_gap
                        })

                    return {
                        "results": materials[:limit],
                        "count": len(materials),
                        "original_query": query,
                        "search_type": "external_search",
                        "formula_used": external_formula
                    }
        
        except Exception as external_error:
            print(f"External search failed: {external_error}")
            # 外部搜索失败，继续尝试其他方法
        
        # 5. 尝试按元素搜索（单元素材料）
        element_symbol = _try_extract_element_symbol(query)
        if element_symbol:
            print(f"Trying element search: '{query}' -> '{element_symbol}'")
            
            with MPRester(API_KEY) as mpr:
                results = mpr.materials.summary.search(
                    elements=[element_symbol],
                    fields=["material_id", "formula_pretty", "symmetry.symbol",
                            "volume", "density", "elements", "band_gap"]
                )
                
                materials = []
                for doc in results:
                    materials.append({
                        "material_id": doc.material_id,
                        "formula_pretty": doc.formula_pretty,
                        "symmetry": doc.symmetry.dict() if hasattr(doc, 'symmetry') and doc.symmetry else None,
                        "volume": doc.volume,
                        "density": doc.density,
                        "elements": doc.elements,
                        "band_gap": doc.band_gap
                    })

                return {
                    "results": materials[:limit],
                    "count": len(materials),
                    "original_query": query,
                    "search_type": "element",
                    "element_used": element_symbol
                }
        
        # 6. 所有方法都失败，返回空结果
        return {
            "results": [],
            "count": 0,
            "original_query": query,
            "search_type": "no_match",
            "message": "No matching materials found. Try using chemical formula or Materials Project ID."
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}")


@app.get("/materials/search/name")
async def search_materials_by_name(query: str, limit: int = 10):
    """
    专门通过材料名称搜索 - 支持中英文名称
    """
    try:
        # 转换为小写进行匹配
        query_lower = query.strip().lower()
        
        # 检查名称映射
        if query_lower in MATERIAL_NAME_MAPPING:
            formula = MATERIAL_NAME_MAPPING[query_lower]
            # 使用转换后的化学式进行搜索
            with MPRester(API_KEY) as mpr:
                results = mpr.materials.summary.search(
                    formula=formula,
                    fields=["material_id", "formula_pretty", "symmetry.symbol",
                            "volume", "density", "elements", "band_gap"]
                )
                # 手动限制结果数量
                if results:
                    results = results[:limit]
                
                materials = []
                for doc in results:
                    materials.append({
                        "material_id": doc.material_id,
                        "formula_pretty": doc.formula_pretty,
                        "symmetry": doc.symmetry.dict() if hasattr(doc, 'symmetry') and doc.symmetry else None,
                        "volume": doc.volume,
                        "density": doc.density,
                        "elements": doc.elements,
                        "band_gap": doc.band_gap
                    })
                
                return {
                    "results": materials,
                    "count": len(materials),
                    "name_query": query,
                    "formula_used": formula
                }
        else:
            return {
                "results": [],
                "count": 0,
                "name_query": query,
                "message": "No matching material name found. Try using chemical formula instead."
            }
            
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Name search failed: {str(e)}")


@app.get("/materials/{material_id}")
async def get_material_details(material_id: str):
    """
    获取材料详细信息
    """
    try:
        with MPRester(API_KEY) as mpr:
            material = mpr.materials.summary.search(
                material_ids=[material_id],
                fields=["material_id", "formula_pretty", "symmetry",
                        "volume", "density", "elements", "band_gap",
                        "energy_per_atom", "formation_energy_per_atom",
                        "nsites", "nelements", "total_magnetization",
                        "is_magnetic", "is_metal", "is_gap_direct",
                        "efermi", "xas", "bulk_modulus", "shear_modulus",
                        "energy_above_hull", "weighted_surface_energy",
                        "weighted_work_function", "universal_anisotropy",
                        "homogeneous_poisson"]
            )

            if not material:
                raise HTTPException(
                    status_code=404, detail="Material not found")

            doc = material[0]
            return {
                "material_id": doc.material_id,
                "formula_pretty": doc.formula_pretty,
                "symmetry": doc.symmetry.dict() if hasattr(doc, 'symmetry') and doc.symmetry else None,
                "volume": doc.volume,
                "density": doc.density,
                "elements": doc.elements,
                "band_gap": doc.band_gap,
                "energy_per_atom": doc.energy_per_atom,
                "formation_energy_per_atom": doc.formation_energy_per_atom,
                "nsites": doc.nsites,
                "nelements": doc.nelements,
                "total_magnetization": getattr(doc, 'total_magnetization', None),
                "is_magnetic": getattr(doc, 'is_magnetic', None),
                "is_metal": getattr(doc, 'is_metal', None),
                "is_gap_direct": getattr(doc, 'is_gap_direct', None),
                "efermi": getattr(doc, 'efermi', None),
                "xas": getattr(doc, 'xas', None),
                "bulk_modulus": getattr(doc, 'bulk_modulus', None),
                "shear_modulus": getattr(doc, 'shear_modulus', None),
                "energy_above_hull": getattr(doc, 'energy_above_hull', None),
                "surface_energy": getattr(doc, 'weighted_surface_energy', None),
                "work_function": getattr(doc, 'weighted_work_function', None),
                "universal_anisotropy": getattr(doc, 'universal_anisotropy', None),
                "poisson_ratio": getattr(doc, 'homogeneous_poisson', None)
            }

    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Failed to fetch material details: {str(e)}")


@app.get("/materials/{material_id}/structure")
async def get_crystal_structure(material_id: str):
    """
    获取晶体结构数据
    """
    try:
        with MPRester(API_KEY) as mpr:
            structure = mpr.get_structure_by_material_id(material_id)

            if not structure:
                raise HTTPException(
                    status_code=404, detail="Structure not found")

            # 转换为可序列化的格式
            lattice = {
                "matrix": structure.lattice.matrix.tolist(),
                "a": structure.lattice.a,
                "b": structure.lattice.b,
                "c": structure.lattice.c,
                "alpha": structure.lattice.alpha,
                "beta": structure.lattice.beta,
                "gamma": structure.lattice.gamma,
                "volume": structure.lattice.volume
            }

            sites = []
            for site in structure:
                sites.append({
                    "species": [str(specie) for specie in site.species],
                    "xyz": site.frac_coords.tolist(),
                    "label": site.label,
                    "properties": site.properties
                })

            return {
                "material_id": material_id,
                "lattice": lattice,
                "sites": sites,
                "charge": structure.charge
            }

    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Failed to fetch crystal structure: {str(e)}")


@app.get("/materials/{material_id}/electronic")
async def get_electronic_structure(material_id: str):
    """
    获取电子结构数据
    """
    try:
        with MPRester(API_KEY) as mpr:
            electronic_data = mpr.materials.electronic_structure.search(
                material_ids=[material_id],
                fields=["bandstructure", "dos", "band_gap"]
            )

            if not electronic_data:
                raise HTTPException(
                    status_code=404, detail="Electronic structure data not found")

            # 简化返回数据，避免过于复杂
            doc = electronic_data[0]
            return {
                "material_id": material_id,
                "band_gap": doc.band_gap.dict() if doc.band_gap else None,
                "has_bandstructure": doc.bandstructure is not None,
                "has_dos": doc.dos is not None
            }

    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Failed to fetch electronic structure: {str(e)}")


@app.get("/examples")
async def get_example_materials():
    """
    获取示例材料列表
    """
    examples = [
        {"material_id": "mp-149", "formula_pretty": "Si", "name": "Silicon"},
        {"material_id": "mp-66", "formula_pretty": "C", "name": "Diamond"},
        {"material_id": "mp-1840", "formula_pretty": "SiO₂", "name": "Quartz"},
        {"material_id": "mp-30", "formula_pretty": "Fe", "name": "Iron"},
        {"material_id": "mp-134", "formula_pretty": "Al", "name": "Aluminum"}
    ]
    return {"examples": examples}

@app.get("/materials/{material_id}/papers", response_model=List[PaperResult])
async def get_material_papers(material_id: str, limit: int = 100):
    """
    获取关于材料的学术论文信息
    """
    try:
        # 首先获取材料信息以获取化学式
        with MPRester(API_KEY) as mpr:
            material = mpr.materials.summary.search(
                material_ids=[material_id],
                fields=["formula_pretty", "elements"]
            )
            
            if not material:
                raise HTTPException(status_code=404, detail="Material not found")
            
            formula = material[0].formula_pretty
            elements = [str(elem) for elem in material[0].elements]  # Convert Element objects to strings
            
            # 尝试使用百度搜索论文
            try:
                if formula == "Si":
                    query = "硅材料研究 半导体 silicon material research"
                else:
                    query = f"{formula} 材料研究 {formula} material research"
                papers = await baidu_searcher.search_papers(query, limit)
                
                # 如果没有结果，尝试更简单的查询
                if not papers:
                    query = f"{formula} {formula} material"
                    papers = await baidu_searcher.search_papers(query, limit)
            except Exception as e:
                print(f"Baidu search failed: {e}")
                papers = []
            
            # 如果百度搜索失败，提供模拟数据
            if not papers:
                papers = _generate_mock_papers(material_id, formula, elements)
            
            return papers[:limit]
            
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Failed to fetch papers: {str(e)}"
        )


@app.get("/materials/{material_id}/prices", response_model=List[PriceResult])
async def get_material_prices(material_id: str, limit: int = 50):
    """
    获取材料的价格和供应商信息
    """
    try:
        # 首先获取材料信息以获取化学式
        with MPRester(API_KEY) as mpr:
            material = mpr.materials.summary.search(
                material_ids=[material_id],
                fields=["formula_pretty", "elements"]
            )
            
            if not material:
                raise HTTPException(status_code=404, detail="Material not found")
            
            formula = material[0].formula_pretty
            elements = [str(elem) for elem in material[0].elements]  # Convert Element objects to strings
            
            # 尝试使用百度搜索价格信息
            try:
                if formula == "Si":
                    query = "硅 价格 供应商 silicon price supplier"
                else:
                    query = f"{formula} 价格 供应商 {formula} price supplier"
                prices = await baidu_searcher.search_prices(query, limit)
                
                # 如果没有结果，尝试更简单的查询
                if not prices:
                    query = f"{formula} price"
                    prices = await baidu_searcher.search_prices(query, limit)
            except Exception as e:
                print(f"Baidu price search failed: {e}")
                prices = []
            
            # 如果百度搜索失败，提供模拟数据
            if not prices:
                prices = _generate_mock_prices(material_id, formula, elements)
            
            return prices[:limit]
            
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Failed to fetch price information: {str(e)}"
        )


@app.get("/search/baidu/papers")
async def search_papers_direct(query: str, limit: int = 100):
    """
    直接搜索学术论文
    """
    try:
        papers = await baidu_searcher.search_papers(query, limit)
        return {"results": papers, "count": len(papers)}
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Failed to search papers: {str(e)}"
        )


@app.get("/search/baidu/prices")
async def search_prices_direct(query: str, limit: int = 50):
    """
    直接搜索价格信息
    """
    try:
        prices = await baidu_searcher.search_prices(query, limit)
        return {"results": prices, "count": len(prices)}
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Failed to search prices: {str(e)}"
        )


def _generate_mock_papers(material_id: str, formula: str, elements: List[str]) -> List[Dict[str, Any]]:
        """生成模拟论文数据"""
        import random
        from datetime import datetime, timedelta
        
        paper_templates = [
            f"{formula}材料的电子结构研究",
            f"{formula}晶体生长与性能分析",
            f"基于{formula}的新型半导体器件",
            f"{formula}纳米材料的制备与应用",
            f"{formula}在光伏领域的应用研究",
            f"{formula}薄膜的物理性质研究",
            f"{' '.join(elements)}化合物材料的最新进展",
            f"{formula}量子点的合成与表征",
            f"{formula}超导材料的研究进展",
            f"{formula}在能源存储中的应用"
        ]
        
        sources = ["scholar.google.com", "arxiv.org", "researchgate.net", "academia.edu", "nature.com", "science.org"]
        
        papers = []
        for i in range(10):
            papers.append({
                "title": random.choice(paper_templates),
                "url": f"https://example.com/paper/{material_id}-{i}",
                "abstract": f"本文研究了{formula}材料的物理化学性质，探讨了其在各个领域的应用潜力。研究发现{formula}具有优异的性能，为新材料开发提供了重要参考。",
                "timestamp": (datetime.now() - timedelta(days=random.randint(1, 365))).isoformat(),
                "source": random.choice(sources)
            })
        
        return papers

def _generate_mock_prices(material_id: str, formula: str, elements: List[str]) -> List[Dict[str, Any]]:
        """生成模拟价格数据"""
        import random
        
        price_templates = [
            f"{formula}高纯度材料供应商",
            f"{formula}工业级原料价格",
            f"{formula}纳米粉末供应商",
            f"{formula}晶体材料采购",
            f"{formula}靶材供应商",
            f"{formula}电子级材料",
            f"{formula}半导体级原料",
            f"{formula}光伏材料价格",
            f"{formula}试剂级供应商",
            f"{formula}高纯金属供应商"
        ]
        
        suppliers = ["阿里巴巴", "中国制造网", "1688.com", "淘宝企业店", "京东工业品", "化工原料网"]
        
        prices = []
        for i in range(8):
            price_val = random.randint(100, 5000)
            prices.append({
                "title": random.choice(price_templates),
                "url": f"https://example.com/price/{material_id}-{i}",
                "abstract": f"提供高纯度{formula}材料，纯度99.99%以上，批量采购优惠。支持定制各种规格，欢迎咨询。",
                "timestamp": datetime.now().isoformat(),
                "source": random.choice(suppliers),
                "price": str(price_val),
                "currency": "CNY",
                "supplier": random.choice(suppliers),
                "quantity": f"{random.choice(['100g', '1kg', '500g', '5kg'])}"
            })
        
        return prices

def _is_chemical_formula(text: str) -> bool:
    """检查字符串是否是化学式格式"""
    # 化学式通常包含字母和数字，可能有括号
    if not any(c.isalpha() for c in text):
        return False
    
    # 检查是否包含常见的化学元素符号模式
    chemical_patterns = [
        r'^[A-Z][a-z]?\d*',  # 单个元素如 H, He, O2
        r'^[A-Z][a-z]?\d*[A-Z][a-z]?\d*',  # 多个元素如 H2O, NaCl
        r'\(.*\)\d*',  # 带括号的化学式如 (NH4)2SO4
        r'\[.*\]\d*',  # 带方括号的化学式
    ]
    
    import re
    for pattern in chemical_patterns:
        if re.match(pattern, text):
            return True
    
    # 检查是否包含常见的化学元素
    common_elements = ['H', 'He', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne',
                      'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'Cl', 'Ar', 'K', 'Ca',
                      'Fe', 'Cu', 'Zn', 'Ag', 'Au', 'Hg', 'Pb']
    
    for element in common_elements:
        if element in text:
            return True
    
    return False

async def _search_formula_from_external(query: str):
    """使用外部搜索获取化学式"""
    try:
        # 使用百度搜索获取可能的化学式
        search_results = await baidu_searcher.search_papers(f"{query} 化学式", 5)
        
        # 从搜索结果中提取化学式
        for result in search_results:
            title = result.get('title', '')
            abstract = result.get('abstract', '')
            
            # 在标题和摘要中查找化学式
            text_to_search = f"{title} {abstract}"
            
            # 查找化学式模式（字母和数字组合）
            import re
            formula_patterns = [
                r'\b[A-Z][a-z]?\d*[A-Z][a-z]?\d*\b',  # 如 H2O, NaCl
                r'\b[A-Z][a-z]?\d*\b',  # 如 O2, H2
                r'\([A-Za-z\d]+\)\d*',  # 如 (NH4)2SO4
            ]
            
            for pattern in formula_patterns:
                matches = re.findall(pattern, text_to_search)
                for match in matches:
                    # 验证找到的字符串确实是化学式
                    if _is_chemical_formula(match):
                        return match
        
        return None
        
    except Exception as e:
        print(f"External formula search failed: {e}")
        return None

def _try_extract_element_symbol(query: str):
    """尝试从查询中提取元素符号"""
    # 元素名称到符号的映射
    element_name_to_symbol = {
        'hydrogen': 'H', '氫': 'H', '氢': 'H',
        'helium': 'He', '氦': 'He',
        'lithium': 'Li', '鋰': 'Li', '锂': 'Li',
        'beryllium': 'Be', '鈹': 'Be', '铍': 'Be',
        'boron': 'B', '硼': 'B',
        'carbon': 'C', '碳': 'C',
        'nitrogen': 'N', '氮': 'N',
        'oxygen': 'O', '氧': 'O',
        'fluorine': 'F', '氟': 'F',
        'neon': 'Ne', '氖': 'Ne',
        'sodium': 'Na', '鈉': 'Na', '钠': 'Na',
        'magnesium': 'Mg', '鎂': 'Mg', '镁': 'Mg',
        'aluminum': 'Al', 'aluminium': 'Al', '鋁': 'Al', '铝': 'Al',
        'silicon': 'Si', '硅': 'Si', '矽': 'Si',
        'phosphorus': 'P', '磷': 'P',
        'sulfur': 'S', '硫': 'S',
        'chlorine': 'Cl', '氯': 'Cl',
        'argon': 'Ar', '氬': 'Ar', '氩': 'Ar',
        'potassium': 'K', '鉀': 'K', '钾': 'K',
        'calcium': 'Ca', '鈣': 'Ca', '钙': 'Ca',
        'iron': 'Fe', '鐵': 'Fe', '铁': 'Fe',
        'copper': 'Cu', '銅': 'Cu', '铜': 'Cu',
        'zinc': 'Zn', '鋅': 'Zn', '锌': 'Zn',
        'silver': 'Ag', '銀': 'Ag', '银': 'Ag',
        'gold': 'Au', '金': 'Au',
        'mercury': 'Hg', '汞': 'Hg',
        'lead': 'Pb', '鉛': 'Pb', '铅': 'Pb'
    }
    
    query_lower = query.lower().strip()
    
    # 直接检查元素名称映射
    if query_lower in element_name_to_symbol:
        return element_name_to_symbol[query_lower]
    
    # 检查是否是单个字母的元素符号
    if len(query) == 1 and query.isalpha() and query.upper() in ['H', 'B', 'C', 'N', 'O', 'F', 'P', 'S', 'K', 'V', 'Y', 'I', 'W', 'U']:
        return query.upper()
    
    # 检查是否是两个字母的元素符号
    if len(query) == 2 and query[0].isalpha() and query[1].isalpha():
        potential_symbol = query[0].upper() + query[1].lower()
        if potential_symbol in ['He', 'Li', 'Be', 'Ne', 'Na', 'Mg', 'Al', 'Si', 'Cl', 'Ar', 'Ca', 'Sc', 'Ti', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ga', 'Ge', 'As', 'Se', 'Br', 'Kr', 'Rb', 'Sr', 'Zr', 'Nb', 'Mo', 'Tc', 'Ru', 'Rh', 'Pd', 'Ag', 'Cd', 'In', 'Sn', 'Sb', 'Te', 'Xe', 'Cs', 'Ba', 'La', 'Ce', 'Pr', 'Nd', 'Pm', 'Sm', 'Eu', 'Gd', 'Tb', 'Dy', 'Ho', 'Er', 'Tm', 'Yb', 'Lu', 'Hf', 'Ta', 'Re', 'Os', 'Ir', 'Pt', 'Au', 'Hg', 'Tl', 'Pb', 'Bi', 'Po', 'At', 'Rn', 'Fr', 'Ra', 'Ac', 'Th', 'Pa', 'Np', 'Pu', 'Am', 'Cm', 'Bk', 'Cf', 'Es', 'Fm', 'Md', 'No', 'Lr']:
            return potential_symbol
    
    return None


async def call_ai_api(prompt: str, temperature: Optional[float] = None, max_tokens: Optional[int] = None):
    """
    调用真正的AI API - 使用配置文件参数
    """
    ai_config = AI_CONFIG["ai_api"]
    
    # 使用配置中的默认值或传入的参数
    ai_url = ai_config["url"]
    model = ai_config["model"]
    timeout = ai_config["timeout"]
    stop_tokens = ai_config["stop_tokens"]
    
    # 如果未提供参数，使用配置中的默认值
    if temperature is None:
        temperature = ai_config["default_temperature"]
    if max_tokens is None:
        max_tokens = ai_config["default_max_tokens"]
    
    # 构建请求数据
    request_data = {
        "model": model,
        "prompt": prompt,
        "stream": True,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "stop": stop_tokens
    }
    
    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.post(
                ai_url,
                json=request_data,
                headers={"Content-Type": "application/json"}
            )
            response.raise_for_status()
            
            return response
            
    except httpx.TimeoutException:
        raise HTTPException(status_code=504, detail="AI API timeout")
    except httpx.HTTPStatusError as e:
        raise HTTPException(status_code=e.response.status_code, detail=f"AI API error: {e}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"AI API connection failed: {str(e)}")


def build_ai_prompt(user_question: str, material_data: Optional[Dict[str, Any]] = None) -> str:
    """
    构建AI提示词 - 使用配置文件模板
    """
    prompt_config = AI_CONFIG["prompt_templates"]
    system_prompt = f"<|im_start|>system\n{prompt_config['system_prompt']}<|im_end|>\n"
    
    if material_data:
        # 准备材料信息参数
        material_params = {
            "material_id": material_data.get('material_id', '未知'),
            "formula_pretty": material_data.get('formula_pretty', '未知'),
            "crystal_system": material_data.get('symmetry', {}).get('crystal_system', '未知') if material_data.get('symmetry') else '未知',
            "space_group": material_data.get('symmetry', {}).get('symbol', '未知') if material_data.get('symmetry') else '未知',
            "band_gap": material_data.get('band_gap', '未知'),
            "density": material_data.get('density', '未知'),
            "volume": material_data.get('volume', '未知')
        }
        
        # 使用模板格式化材料信息
        material_info = prompt_config["material_info_template"].format(**material_params)
    else:
        material_info = "当前没有选择特定材料。"
    
    user_prompt = f"{system_prompt}\n<|im_start|>user\n{material_info}\n用户问题：{user_question}<|im_end|>\n"
    
    return user_prompt


from fastapi.responses import StreamingResponse

@app.post("/ai/chat")
async def ai_chat_stream(request: AIRequest):
    """
    AI对话接口 - 流式响应
    """
    async def generate():
        try:
            # 构建提示词
            prompt = build_ai_prompt(request.prompt, request.material_data)
            
            # 调用AI API
            response = await call_ai_api(
                prompt=prompt,
                temperature=request.temperature,
                max_tokens=request.max_tokens
            )
            
            # 处理流式响应
            buffer = ""
            async for chunk in response.aiter_text():
                if chunk.strip():
                    buffer += chunk
                    
                    # 按行分割处理
                    lines = buffer.split('\n')
                    buffer = lines.pop() if lines else ""  # 保留最后一行（可能不完整）
                    
                    for line in lines:
                        line = line.strip()
                        if line and line.startswith("data: "):
                            # 提取JSON部分
                            json_str = line[6:].strip()  # 移除"data: "前缀
                            
                            # 检查是否是结束标记
                            if json_str == "[DONE]":
                                print("收到AI API结束标记")
                                continue
                            
                            # 尝试解析JSON
                            try:
                                data = json.loads(json_str)
                                
                                # 提取文本内容
                                content = ""
                                if data.get("choices") and isinstance(data["choices"], list) and len(data["choices"]) > 0:
                                    choice = data["choices"][0]
                                    if choice.get("text") is not None:
                                        content = choice["text"]
                                
                                # 如果text字段为空，检查其他可能的字段
                                if not content:
                                    if data.get("content"):
                                        content = data["content"]
                                    elif data.get("text"):
                                        content = data["text"]
                                    elif data.get("response"):
                                        content = data["response"]
                                
                                # 只发送文本内容，不发送整个JSON对象
                                if content:
                                    # 为了更好的流式效果，逐字符发送
                                    for char in content:
                                        yield "data: " + json.dumps({"content": char, "is_final": False}) + "\n\n"
                                        # 添加小延迟以创建更好的流式效果
                                        await asyncio.sleep(0.01)
                                    
                            except json.JSONDecodeError as e:
                                print(f"JSON解析错误: {e}, 原始数据: {json_str}")
                                # 如果不是JSON格式，尝试直接提取文本内容
                                # 检查是否是包含"text"字段的字符串
                                if '"text":' in json_str:
                                    try:
                                        # 尝试提取text字段的值
                                        import re
                                        text_match = re.search(r'"text":\s*"([^"]*)"', json_str)
                                        if text_match:
                                            content = text_match.group(1)
                                            for char in content:
                                                yield "data: " + json.dumps({"content": char, "is_final": False}) + "\n\n"
                                                await asyncio.sleep(0.01)
                                        else:
                                            # 直接返回原始文本
                                            for char in json_str:
                                                yield "data: " + json.dumps({"content": char, "is_final": False}) + "\n\n"
                                                await asyncio.sleep(0.01)
                                    except:
                                        for char in json_str:
                                            yield "data: " + json.dumps({"content": char, "is_final": False}) + "\n\n"
                                            await asyncio.sleep(0.01)
                                else:
                                    for char in json_str:
                                        yield "data: " + json.dumps({"content": char, "is_final": False}) + "\n\n"
                                        await asyncio.sleep(0.01)
            
            # 处理剩余的数据
            if buffer.strip() and buffer.startswith("data: "):
                try:
                    json_str = buffer[6:].strip()
                    data = json.loads(json_str)
                    
                    content = ""
                    if data.get("choices") and isinstance(data["choices"], list) and len(data["choices"]) > 0:
                        choice = data["choices"][0]
                        if choice.get("text") is not None:
                            content = choice["text"]
                    
                    if not content:
                        if data.get("content"):
                            content = data["content"]
                        elif data.get("text"):
                            content = data["text"]
                        elif data.get("response"):
                            content = data["response"]
                    
                    if content:
                        for char in content:
                            yield "data: " + json.dumps({"content": char, "is_final": False}) + "\n\n"
                            await asyncio.sleep(0.01)
                except json.JSONDecodeError as e:
                    print(f"剩余数据JSON解析错误: {e}")
                    for char in buffer[6:].strip():
                        yield "data: " + json.dumps({"content": char, "is_final": False}) + "\n\n"
                        await asyncio.sleep(0.01)
            
            # 发送结束标记
            yield "data: " + json.dumps({"content": "", "is_final": True}) + "\n\n"
            
        except Exception as e:
            yield "data: " + json.dumps({"content": f'AI对话出错：{str(e)}', "is_final": True}) + "\n\n"
    
    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, OPTIONS",
            "Access-Control-Allow-Headers": "Content-Type, Authorization, X-Requested-With",
            "Content-Type": "text/event-stream; charset=utf-8"
        }
    )


@app.post("/ai/chat/simple")
async def ai_chat_simple(request: AIRequest):
    """
    AI对话接口 - 简单响应（非流式）
    """
    try:
        # 构建提示词
        prompt = build_ai_prompt(request.prompt, request.material_data)
        
        # 调用AI API
        response = await call_ai_api(
            prompt=prompt,
            temperature=request.temperature,
            max_tokens=request.max_tokens
        )
        
        # 收集所有响应内容
        full_response = ""
        async for chunk in response.aiter_text():
            if chunk.strip():
                # 检查是否是SSE格式（以"data: "开头）
                if chunk.startswith("data: "):
                    # 提取JSON部分
                    json_str = chunk[6:].strip()  # 移除"data: "前缀
                    try:
                        data = json.loads(json_str)
                        
                        # 尝试从不同的字段获取内容
                        content = ""
                        if data.get("choices") and data["choices"] and data["choices"][0]:
                            choice = data["choices"][0]
                            if choice.get("text"):
                                content = choice["text"]
                            elif choice.get("message") and choice["message"].get("content"):
                                content = choice["message"]["content"]
                            elif choice.get("delta") and choice["delta"].get("content"):
                                content = choice["delta"]["content"]
                        elif data.get("content"):
                            content = data["content"]
                        elif data.get("text"):
                            content = data["text"]
                        elif data.get("response"):
                            content = data["response"]
                        
                        full_response += content
                    except json.JSONDecodeError:
                        full_response += json_str
                else:
                    # 如果不是SSE格式，尝试直接解析为JSON
                    try:
                        data = json.loads(chunk)
                        
                        # 尝试从不同的字段获取内容
                        content = ""
                        if data.get("choices") and data["choices"] and data["choices"][0]:
                            choice = data["choices"][0]
                            if choice.get("text"):
                                content = choice["text"]
                            elif choice.get("message") and choice["message"].get("content"):
                                content = choice["message"]["content"]
                            elif choice.get("delta") and choice["delta"].get("content"):
                                content = choice["delta"]["content"]
                        elif data.get("content"):
                            content = data["content"]
                        elif data.get("text"):
                            content = data["text"]
                        elif data.get("response"):
                            content = data["response"]
                        
                        full_response += content
                    except json.JSONDecodeError:
                        full_response += chunk
        
        return {"response": full_response}
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"AI对话失败：{str(e)}")


if __name__ == "__main__":
    import uvicorn
    import argparse
    
    parser = argparse.ArgumentParser(description='Molecule Scope API Server')
    parser.add_argument('--port', type=int, default=8001, help='Port to run the server on')
    parser.add_argument('--host', type=str, default='0.0.0.0', help='Host to bind the server to')
    
    args = parser.parse_args()
    
    uvicorn.run(app, host=args.host, port=args.port)
